import torch
import torch.nn as nn
import torch.nn.functional as F
from .base_module import BaseModule

from .batchnorm import SynchronizedBatchNorm2d
from .ultrasound_rfb import UltrasoundRFB, AdaptiveRFB
from .pvtv2 import pvt_v2_b2
from .Res2Net_v1b import res2net50_v1b_26w_4s

BatchNorm2d = SynchronizedBatchNorm2d

def conv3x3_bn_relu(in_planes, out_planes, stride=1):
    "3x3 convolution + BN + relu"
    return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=3,
                      stride=stride, padding=1, bias=False),
            BatchNorm2d(out_planes),
            nn.ReLU(inplace=True),
            )

class BasicConv2d(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv2d(in_planes, out_planes,
                              kernel_size=kernel_size, stride=stride,
                              padding=padding, dilation=dilation, bias=False)
        self.bn = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x

class UltrasoundPreprocessing(nn.Module):
    """超声图像预处理模块"""
    def __init__(self, in_channels=1, out_channels=3):
        super(UltrasoundPreprocessing, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.conv3 = nn.Conv2d(32, out_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.bn2 = nn.BatchNorm2d(32)
        self.bn3 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.relu(self.bn3(self.conv3(x)))
        return x

class EdgeAwareAttention(nn.Module):
    """边缘感知注意力模块"""
    def __init__(self, in_channels):
        super(EdgeAwareAttention, self).__init__()
        # 使用分组卷积进行边缘检测，每个通道独立处理
        self.edge_conv = nn.Conv2d(in_channels, in_channels, 3, padding=1, 
                                  groups=in_channels, bias=False)
        self.attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, in_channels // 4, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels // 4, in_channels, 1),
            nn.Sigmoid()
        )
        
        # 初始化边缘检测核
        edge_kernel = torch.tensor([
            [-1, -1, -1],
            [-1,  8, -1],
            [-1, -1, -1]
        ], dtype=torch.float32)
        
        # 为每个通道设置相同的边缘检测核
        weight = edge_kernel.unsqueeze(0).unsqueeze(0).repeat(in_channels, 1, 1, 1)
        self.edge_conv.weight.data = weight
        
    def forward(self, x):
        edge = self.edge_conv(x)
        att = self.attention(x)
        return x * (1 + 0.1 * edge) * att

class ImprovedAggregation(nn.Module):
    """改进的特征聚合模块"""
    def __init__(self, channel, num_class):
        super(ImprovedAggregation, self).__init__()
        self.relu = nn.ReLU(True)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        
        # 改进的上采样卷积
        self.conv_upsample1 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample2 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample3 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample4 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample5 = BasicConv2d(2*channel, 2*channel, 3, padding=1)
        
        # 注意力机制
        self.attention2 = EdgeAwareAttention(2*channel)
        self.attention3 = EdgeAwareAttention(3*channel)
        
        self.conv_concat2 = BasicConv2d(2*channel, 2*channel, 3, padding=1)
        self.conv_concat3 = BasicConv2d(3*channel, 3*channel, 3, padding=1)
        self.conv4 = BasicConv2d(3*channel, 3*channel, 3, padding=1)
        self.conv5_fg = nn.Conv2d(3*channel, num_class, 1)
        self.conv5_bg = nn.Conv2d(3*channel, num_class, 1)

    def forward(self, x1, x2, x3):
        x1_1 = x1
        x1_up = self.upsample(x1)
        x1_up_up = self.upsample(x1_up)
        x2_up = self.upsample(x2)
        x2_1 = self.conv_upsample1(x1_up) * x2
        x3_1 = self.conv_upsample2(x1_up_up) \
               * self.conv_upsample3(x2_up) * x3

        x2_2 = torch.cat((x2_1, self.conv_upsample4(x1_up)), 1)
        x2_2 = self.attention2(self.conv_concat2(x2_2))

        x3_2 = torch.cat((x3_1, self.conv_upsample5(self.upsample(x2_2))), 1)
        x3_2 = self.attention3(self.conv_concat3(x3_2))

        x = self.conv4(x3_2)
        x_fg = self.conv5_fg(x)
        x_bg = self.conv5_bg(x)

        return x_fg, x_bg

class AEAPraNet_PVT(BaseModule):
    """改进的PRANet V2模型，专门针对超声图像优化"""
    
    def __init__(self, channel=32, num_classes=3, sem_downsample=1, use_softmax=True, 
                 pretrained_model_path=None):
        super(AEAPraNet_PVT, self).__init__(num_classes)
        self.num_classes = num_classes
        self.sem_downsample = sem_downsample
        self.use_softmax = use_softmax

        if pretrained_model_path is not None:
            pretrained_model_path = f'{pretrained_model_path}/pvt_v2_b2.pth'
        
        # Backbone
        self.backbone = pvt_v2_b2()
        if pretrained_model_path is not None:
            save_model = torch.load(pretrained_model_path)
            model_dict = self.backbone.state_dict()
            state_dict = {k: v for k, v in save_model.items() if k in model_dict.keys()}
            model_dict.update(state_dict)
            self.backbone.load_state_dict(model_dict)
        
        # 改进的RFB模块
    
        self.rfb2_1 = AdaptiveRFB(128, channel)
        self.rfb3_1 = AdaptiveRFB(320, channel)
        self.rfb4_1 = AdaptiveRFB(512, channel)
        
        # 改进的特征聚合
        self.agg1 = ImprovedAggregation(channel, self.num_classes)
        
        # DSRA模块（保持原有结构，但添加边缘感知）
        self.ra4_conv1 = BasicConv2d(512, 256, kernel_size=1)
        self.ra4_conv2 = BasicConv2d(256, 256, kernel_size=5, padding=2)
        self.ra4_conv3 = BasicConv2d(256, 256, kernel_size=5, padding=2)
        self.ra4_conv4 = BasicConv2d(256, 256, kernel_size=5, padding=2)
        self.ra4_conv5_fg = BasicConv2d(256, num_classes, kernel_size=1)
        self.ra4_conv5_bg = BasicConv2d(256, num_classes, kernel_size=1)
        
        self.ra3_conv1 = BasicConv2d(320, 64, kernel_size=1)
        self.ra3_conv2 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra3_conv3 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra3_conv4_fg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)
        self.ra3_conv4_bg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)
        
        self.ra2_conv1 = BasicConv2d(128, 64, kernel_size=1)
        self.ra2_conv2 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra2_conv3 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra2_conv4_fg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)
        self.ra2_conv4_bg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)

    def forward(self, x):        
        # 编码器阶段
        x1, x2, x3, x4 = self.backbone(x)
        
        # 改进的RFB处理
        x2_rfb = self.rfb2_1(x2)
        x3_rfb = self.rfb3_1(x3)
        x4_rfb = self.rfb4_1(x4)
        
        # 特征融合
        ra5_feat_fg, ra5_feat_bg = self.agg1(x4_rfb, x3_rfb, x2_rfb)
        lateral_map_5_fg = F.interpolate(ra5_feat_fg, scale_factor=8/self.sem_downsample, mode='bilinear')
        lateral_map_5_bg = F.interpolate(ra5_feat_bg, scale_factor=8/self.sem_downsample, mode='bilinear')
        
        # DSRA3
        crop_4_fg = F.interpolate(ra5_feat_fg, scale_factor=0.25, mode='bilinear')
        crop_4_bg = F.interpolate(ra5_feat_bg, scale_factor=0.25, mode='bilinear')
        
        x = self.ra4_conv1(x4)
        x = F.relu(self.ra4_conv2(x))
        x = F.relu(self.ra4_conv3(x))
        x = F.relu(self.ra4_conv4(x))
        
        ra4_feat_fg = self.ra4_conv5_fg(x)
        ra4_feat_bg = self.ra4_conv5_bg(x)
        
        if self.use_softmax:
            ra4_feat_fg = ra4_feat_fg + ra4_feat_fg.mul(F.softmax(crop_4_fg-crop_4_bg, dim=1))
        else:
            ra4_feat_fg = ra4_feat_fg + ra4_feat_fg.mul(crop_4_fg-crop_4_bg)
        
        lateral_map_4_fg = F.interpolate(ra4_feat_fg, scale_factor=32/self.sem_downsample, mode='bilinear')
        lateral_map_4_bg = F.interpolate(ra4_feat_bg, scale_factor=32/self.sem_downsample, mode='bilinear')
        
        # DSRA2
        crop_3_fg = F.interpolate(ra4_feat_fg, scale_factor=2, mode='bilinear')
        crop_3_bg = F.interpolate(ra4_feat_bg, scale_factor=2, mode='bilinear')
        
        x = self.ra3_conv1(x3)
        x = F.relu(self.ra3_conv2(x))
        x = F.relu(self.ra3_conv3(x))
        
        ra3_feat_fg = self.ra3_conv4_fg(x)
        ra3_feat_bg = self.ra3_conv4_bg(x)
        
        if self.use_softmax:
            ra3_feat_fg = ra3_feat_fg + ra3_feat_fg.mul(F.softmax(crop_3_fg-crop_3_bg, dim=1))
        else:
            ra3_feat_fg = ra3_feat_fg + ra3_feat_fg.mul(crop_3_fg-crop_3_bg)
        
        lateral_map_3_fg = F.interpolate(ra3_feat_fg, scale_factor=16/self.sem_downsample, mode='bilinear')
        lateral_map_3_bg = F.interpolate(ra3_feat_bg, scale_factor=16/self.sem_downsample, mode='bilinear')
        
        # DSRA1
        crop_2_fg = F.interpolate(ra3_feat_fg, scale_factor=2, mode='bilinear')
        crop_2_bg = F.interpolate(ra3_feat_bg, scale_factor=2, mode='bilinear')
        
        x = self.ra2_conv1(x2)
        x = F.relu(self.ra2_conv2(x))
        x = F.relu(self.ra2_conv3(x))
        
        ra2_feat_fg = self.ra2_conv4_fg(x)
        ra2_feat_bg = self.ra2_conv4_bg(x)
        
        if self.use_softmax:
            ra2_feat_fg = ra2_feat_fg + ra2_feat_fg.mul(F.softmax(crop_2_fg-crop_2_bg, dim=1))
        else:
            ra2_feat_fg = ra2_feat_fg + ra2_feat_fg.mul(crop_2_fg-crop_2_bg)
        
        lateral_map_2_fg = F.interpolate(ra2_feat_fg, scale_factor=8/self.sem_downsample, mode='bilinear')
        lateral_map_2_bg = F.interpolate(ra2_feat_bg, scale_factor=8/self.sem_downsample, mode='bilinear')
        
        return (lateral_map_2_fg, lateral_map_3_fg, lateral_map_4_fg, lateral_map_5_fg,
                lateral_map_2_bg, lateral_map_3_bg, lateral_map_4_bg, lateral_map_5_bg)

class AEAPraNet_ResNet(BaseModule):
    """改进的ResNet-based PRANet V2模型，专门针对超声图像优化"""
    
    def __init__(self, channel=32, num_classes=3, sem_downsample=1, use_softmax=True, 
                 pretrained_model_path=None, use_adaptive_rfb=True):
        super(AEAPraNet_ResNet, self).__init__(num_classes)
        self.num_classes = num_classes
        self.sem_downsample = sem_downsample
        self.use_softmax = use_softmax
        self.use_adaptive_rfb = use_adaptive_rfb

        if pretrained_model_path is not None:
            pretrained_model_path = f'{pretrained_model_path}/res2net50_v1b_26w_4s-3cf99910.pth'
        
        # ResNet backbone
        self.backbone = res2net50_v1b_26w_4s(pretrained_model_path=pretrained_model_path)

        # 改进的RFB模块
        if use_adaptive_rfb:
            self.rfb2_1 = AdaptiveRFB(512, channel)
            self.rfb3_1 = AdaptiveRFB(1024, channel)
            self.rfb4_1 = AdaptiveRFB(2048, channel)
        else:
            self.rfb2_1 = UltrasoundRFB(512, channel)
            self.rfb3_1 = UltrasoundRFB(1024, channel)
            self.rfb4_1 = UltrasoundRFB(2048, channel)
        
        # 改进的特征聚合
        self.agg1 = ImprovedAggregation(channel, self.num_classes)
        
        # DSRA模块（保持原有结构，但添加边缘感知）
        self.ra4_conv1 = BasicConv2d(2048, 256, kernel_size=1)
        self.ra4_conv2 = BasicConv2d(256, 256, kernel_size=5, padding=2)
        self.ra4_conv3 = BasicConv2d(256, 256, kernel_size=5, padding=2)
        self.ra4_conv4 = BasicConv2d(256, 256, kernel_size=5, padding=2)
        self.ra4_conv5_fg = BasicConv2d(256, num_classes, kernel_size=1)
        self.ra4_conv5_bg = BasicConv2d(256, num_classes, kernel_size=1)
        
        self.ra3_conv1 = BasicConv2d(1024, 64, kernel_size=1)
        self.ra3_conv2 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra3_conv3 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra3_conv4_fg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)
        self.ra3_conv4_bg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)
        
        self.ra2_conv1 = BasicConv2d(512, 64, kernel_size=1)
        self.ra2_conv2 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra2_conv3 = BasicConv2d(64, 64, kernel_size=3, padding=1)
        self.ra2_conv4_fg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)
        self.ra2_conv4_bg = BasicConv2d(64, num_classes, kernel_size=3, padding=1)

    def forward(self, x, segSize=None):
        # 编码器阶段
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)      # bs, 64, 88, 88
        x1 = self.backbone.layer1(x)      # bs, 256, 88, 88
        x2 = self.backbone.layer2(x1)     # bs, 512, 44, 44
        x3 = self.backbone.layer3(x2)     # bs, 1024, 22, 22
        x4 = self.backbone.layer4(x3)     # bs, 2048, 11, 11

        # 改进的RFB处理
        x2_rfb = self.rfb2_1(x2)        # channel -> 32
        x3_rfb = self.rfb3_1(x3)        # channel -> 32
        x4_rfb = self.rfb4_1(x4)        # channel -> 32

        # 特征融合
        ra5_feat_fg, ra5_feat_bg = self.agg1(x4_rfb, x3_rfb, x2_rfb)
        lateral_map_5_fg = F.interpolate(ra5_feat_fg, scale_factor=8/self.sem_downsample, mode='bilinear')
        lateral_map_5_bg = F.interpolate(ra5_feat_bg, scale_factor=8/self.sem_downsample, mode='bilinear')

        # DSRA3
        crop_4_fg = F.interpolate(ra5_feat_fg, scale_factor=0.25, mode='bilinear')
        crop_4_bg = F.interpolate(ra5_feat_bg, scale_factor=0.25, mode='bilinear')
        
        x = self.ra4_conv1(x4)
        x = F.relu(self.ra4_conv2(x))
        x = F.relu(self.ra4_conv3(x))
        x = F.relu(self.ra4_conv4(x))

        ra4_feat_fg = self.ra4_conv5_fg(x)
        ra4_feat_bg = self.ra4_conv5_bg(x)

        if self.use_softmax:
            ra4_feat_fg = ra4_feat_fg + ra4_feat_fg.mul(F.softmax(crop_4_fg-crop_4_bg, dim=1))
        else:
            ra4_feat_fg = ra4_feat_fg + ra4_feat_fg.mul(crop_4_fg-crop_4_bg)

        lateral_map_4_fg = F.interpolate(ra4_feat_fg, scale_factor=32/self.sem_downsample, mode='bilinear')
        lateral_map_4_bg = F.interpolate(ra4_feat_bg, scale_factor=32/self.sem_downsample, mode='bilinear')

        # DSRA2
        crop_3_fg = F.interpolate(ra4_feat_fg, scale_factor=2, mode='bilinear')
        crop_3_bg = F.interpolate(ra4_feat_bg, scale_factor=2, mode='bilinear')

        x = self.ra3_conv1(x3)
        x = F.relu(self.ra3_conv2(x))
        x = F.relu(self.ra3_conv3(x))

        ra3_feat_fg = self.ra3_conv4_fg(x)
        ra3_feat_bg = self.ra3_conv4_bg(x)

        if self.use_softmax:
            ra3_feat_fg = ra3_feat_fg + ra3_feat_fg.mul(F.softmax(crop_3_fg-crop_3_bg, dim=1))
        else:
            ra3_feat_fg = ra3_feat_fg + ra3_feat_fg.mul(crop_3_fg-crop_3_bg)

        lateral_map_3_fg = F.interpolate(ra3_feat_fg, scale_factor=16/self.sem_downsample, mode='bilinear')
        lateral_map_3_bg = F.interpolate(ra3_feat_bg, scale_factor=16/self.sem_downsample, mode='bilinear')

        # DSRA1
        crop_2_fg = F.interpolate(ra3_feat_fg, scale_factor=2, mode='bilinear')
        crop_2_bg = F.interpolate(ra3_feat_bg, scale_factor=2, mode='bilinear')

        x = self.ra2_conv1(x2)
        x = F.relu(self.ra2_conv2(x))
        x = F.relu(self.ra2_conv3(x))

        ra2_feat_fg = self.ra2_conv4_fg(x)
        ra2_feat_bg = self.ra2_conv4_bg(x)

        if self.use_softmax:
            ra2_feat_fg = ra2_feat_fg + ra2_feat_fg.mul(F.softmax(crop_2_fg-crop_2_bg, dim=1))
        else:
            ra2_feat_fg = ra2_feat_fg + ra2_feat_fg.mul(crop_2_fg-crop_2_bg)

        lateral_map_2_fg = F.interpolate(ra2_feat_fg, scale_factor=8/self.sem_downsample, mode='bilinear')
        lateral_map_2_bg = F.interpolate(ra2_feat_bg, scale_factor=8/self.sem_downsample, mode='bilinear')

        return (lateral_map_2_fg, lateral_map_3_fg, lateral_map_4_fg, lateral_map_5_fg,
                lateral_map_2_bg, lateral_map_3_bg, lateral_map_4_bg, lateral_map_5_bg)

